
#include "DBCompressedStream.h"

#ifdef _DB_USE_ZLIB

#include "DBMemoryBlock.h"

#include <stdio.h>
#include <string.h>
#include <assert.h>
#include "zlib.h"

namespace db
{
    
// ----------------------------------- WriteCompressedStream ----------------------------------------

MemoryBlock* WriteCompressedStream::allocBlock(size_t sizeOf /* = DATA_CHUNK_SIZE */)
{
    MemoryBlock *block = new MemoryBlock(sizeOf);
    m_buffers.push_back(block);
    return block;
}

MemoryBlock* WriteCompressedStream::topBlock() 
{ 
    return m_buffers.back(); 
}


WriteCompressedStream::WriteCompressedStream(WriteStream& outStream) :
    m_outStream(outStream)
{
    allocBlock();
}

bool WriteCompressedStream::write(const void* data, const unsigned int size)
{
    MemoryBlock* b_data = topBlock();
    if (size > DATA_CHUNK_SIZE)
    {
        size_t left_size = b_data->size - b_data->position;

        b_data->write(data, left_size);
        char* c_data = (char*)data + left_size;
        left_size = size - left_size;

        b_data = allocBlock(left_size < DATA_CHUNK_SIZE ? DATA_CHUNK_SIZE : left_size);
        b_data->write(c_data, left_size);
    } 
    else if (b_data->position + size <= b_data->size)
    {
        b_data->write(data, size);
    }
    else
    {
        size_t left_size = b_data->size - b_data->position;

        b_data->write(data, left_size);
        char* c_data = (char*)data + left_size;
        left_size = size - left_size;

        b_data = allocBlock(left_size < DATA_CHUNK_SIZE ? DATA_CHUNK_SIZE : left_size);
        b_data->write(c_data, left_size);
    }
    return true;
}

void WriteCompressedStream::complete()
{
    uLongf totalSize = 0;
    for (size_t i = 0, t = m_buffers.size(); i < t; ++i)
    {
        totalSize += m_buffers[i]->position;
    }
    char* data = (char*)malloc(totalSize);
    char* i_data = data;

    for (size_t i = 0, t = m_buffers.size(); i < t; ++i)
    {
        memcpy(i_data, m_buffers[i]->data, m_buffers[i]->position);
        i_data += m_buffers[i]->size;
    }

    for (size_t i = 0, t = m_buffers.size(); i < t; ++i) delete m_buffers[i];
    m_buffers.clear();

    uLong buffer_size = compressBound(totalSize);
    char* buffer_data = (char*)malloc(buffer_size);

    compress2((Bytef*)buffer_data, &buffer_size, (Bytef*)data, totalSize, Z_DEFAULT_COMPRESSION);

    m_outStream.write(&buffer_size, sizeof(buffer_size));
    m_outStream.write(&totalSize, sizeof(totalSize));
    m_outStream.write(buffer_data, buffer_size);
    m_outStream.complete();

    free(data);
    free(buffer_data);
}

void WriteCompressedStream::prepare()
{
    //
}

// ----------------------------------- ReadCompressedStream ----------------------------------------

ReadCompressedStream::ReadCompressedStream(ReadStream& inStream) :
        m_inStream(inStream)
{
    uLong packed_buffer_size;
    m_inStream.read(&packed_buffer_size, sizeof(packed_buffer_size));
    m_inStream.read(&m_bufferSize, sizeof(m_bufferSize));

    char *packed = (char*)malloc(packed_buffer_size);
    m_inStream.read(packed, packed_buffer_size);

    m_buffer = (char*)malloc(m_bufferSize);
    uncompress((Bytef*)m_buffer, &m_bufferSize, (Bytef*)packed, packed_buffer_size);     

    m_bufferPos = 0;

    free(packed);
}

bool ReadCompressedStream::read(void* data, const unsigned int size)
{
    if (m_bufferPos + size <= m_bufferSize)
    {
        memcpy(data, m_buffer + m_bufferPos, size);
        m_bufferPos += size;
        return true;
    }
    return false;
}

void ReadCompressedStream::complete()
{
    free(m_buffer);
    m_bufferSize = 0;
    m_bufferPos = 0;
}

};

#endif